import numpy as np
import time

def bev_pool(depth, feat, ranks_depth, ranks_feat, ranks_bev,
                       bev_feat_shape, interval_starts, interval_lengths):
    """
    优化后的 BEV Pool 算子，使用 NumPy 张量操作实现（只保留 batch 循环）。

    Args:
        depth: (B, N, D, fH, fW)
        feat: (B, N, fH, fW, C)
        ranks_depth: (N_points,)
        ranks_feat: (N_points,)
        ranks_bev: (N_points,)
        bev_feat_shape: (B, D_Z, D_Y, D_X, C)
        interval_starts, interval_lengths: 暂未使用（如需稀疏优化可引入）

    Returns:
        BEV特征: (B, C, D_Z, D_Y, D_X)
    """
    B, D_Z, D_Y, D_X, C = bev_feat_shape
    bev_feat = np.zeros(bev_feat_shape, dtype=feat.dtype)

    for b in range(B):
        # 1. 找出属于该 batch 的点
        batch_ids_bev = np.unravel_index(ranks_bev, (B, D_Z, D_Y, D_X, C))[0]
        mask_b = (batch_ids_bev == b)
        if not np.any(mask_b):
            continue

        # 2. 提取该 batch 的索引
        r_depth_b = ranks_depth[mask_b]
        r_feat_b = ranks_feat[mask_b]
        r_bev_b = ranks_bev[mask_b]

        # 3. 展平索引 → 多维索引
        b_depth, n, d, fh, fw = np.unravel_index(r_depth_b, depth.shape)
        b_feat, n2, fh2, fw2, c = np.unravel_index(r_feat_b, feat.shape)
        b_bev, dz, dy, dx, c2 = np.unravel_index(r_bev_b, bev_feat_shape)

        # 4. 特征加权
        depth_vals = depth[b, n, d, fh, fw]           # (N_valid,)
        feat_vals = feat[b, n2, fh2, fw2, c]           # (N_valid,)
        weighted_feat = depth_vals * feat_vals        # (N_valid,)

        # 5. 聚合到 BEV 网格（模拟 scatter_add）
        np.add.at(bev_feat[b], (dz, dy, dx, c2), weighted_feat)

    # 6. 转为 (B, C, D_Z, D_Y, D_X)
    return np.transpose(bev_feat, (0, 4, 1, 2, 3))


def bev_pool_optimized(depth, feat, ranks_depth, ranks_feat, ranks_bev,
                       bev_feat_shape, interval_starts, interval_lengths):
    """
    Vectorized version of BEV pooling.

    Args:
        depth: (B, N, D, fH, fW)
        feat:  (B, N, fH, fW, C)
        ranks_depth: (N_points,)
        ranks_feat: (N_points,)
        ranks_bev: (N_points,)
        bev_feat_shape: (B, D_Z, D_Y, D_X, C)
        interval_starts: (N_pillar,)
        interval_lengths: (N_pillar,)

    Returns:
        x: BEV feature in shape (B, C, D_Z, D_Y, D_X)
    """
    B, D_Z, D_Y, D_X, C = bev_feat_shape
    bev_feat = np.zeros(bev_feat_shape, dtype=feat.dtype)

    # Decode the indices
    b_d, n_d, d_d, fh_d, fw_d = np.unravel_index(ranks_depth, depth.shape)
    b_f, n_f, fh_f, fw_f, c_f = np.unravel_index(ranks_feat, feat.shape)
    b_b, dz_b, dy_b, dx_b, c_b = np.unravel_index(ranks_bev, bev_feat_shape)

    # (Optional) consistency check
    assert np.all(b_d == b_f) and np.all(n_d == n_f)
    assert np.all(fh_d == fh_f) and np.all(fw_d == fw_f)
    assert np.all(b_d == b_b) and np.all(c_f == c_b)

    # Compute values to accumulate
    values = depth[b_d, n_d, d_d, fh_d, fw_d] * feat[b_f, n_f, fh_f, fw_f, c_f]

    # Compute linear indices to accumulate into
    linear_indices = np.ravel_multi_index((b_b, dz_b, dy_b, dx_b, c_b), bev_feat_shape)

    # Flatten and scatter-add the values
    flat_bev_feat = bev_feat.reshape(-1)
    np.add.at(flat_bev_feat, linear_indices, values)

    # Reshape and transpose to (B, C, D_Z, D_Y, D_X)
    bev_feat = flat_bev_feat.reshape(bev_feat_shape)
    x = np.transpose(bev_feat, (0, 4, 1, 2, 3))
    return x


# 示例代码
if __name__ == "__main__":
    # 输入参数形状
    B = 1
    N = 2
    D = 3
    fH = 4
    fW = 5
    C = 6
    D_Z = 3
    D_Y = 4
    D_X = 5
    N_points = 1000
    N_pillar = 50

    # 构造输入数据
    depth = np.random.randn(B, N, D, fH, fW)
    feat = np.random.randn(B, N, fH, fW, C)

    # 生成有效索引，确保共享相同 (b, n, fh, fw)
    b_indices = np.random.randint(0, B, N_points)
    n_indices = np.random.randint(0, N, N_points)
    fh_indices = np.random.randint(0, fH, N_points)
    fw_indices = np.random.randint(0, fW, N_points)
    d_indices = np.random.randint(0, D, N_points)
    c_indices = np.random.randint(0, C, N_points)

    ranks_depth = np.ravel_multi_index((b_indices, n_indices, d_indices, fh_indices, fw_indices), depth.shape)
    ranks_feat = np.ravel_multi_index((b_indices, n_indices, fh_indices, fw_indices, c_indices), feat.shape)
    dz_indices = np.random.randint(0, D_Z, N_points)
    dy_indices = np.random.randint(0, D_Y, N_points)
    dx_indices = np.random.randint(0, D_X, N_points)
    ranks_bev = np.ravel_multi_index((b_indices, dz_indices, dy_indices, dx_indices, c_indices),
                                     (B, D_Z, D_Y, D_X, C))

    bev_feat_shape = (B, D_Z, D_Y, D_X, C)

    interval_starts = np.linspace(0, N_points - N_points // N_pillar, N_pillar, dtype=int)
    interval_lengths = np.full(N_pillar, N_points // N_pillar, dtype=int)
    interval_lengths[-1] = N_points - interval_lengths[:-1].sum()  # ensure sum matches N_points

    start_time = time.time()
#    result = bev_pool_optimized(depth, feat, ranks_depth, ranks_feat, ranks_bev,
#                                bev_feat_shape, interval_starts, interval_lengths)
    result = bev_pool(depth, feat, ranks_depth, ranks_feat, ranks_bev,
                                bev_feat_shape, interval_starts, interval_lengths)

    end_time = time.time()

    print(f"BEV Pool 执行时间: {end_time - start_time:.6f} 秒")
    print("输出 BEV 特征形状:", result.shape)

